'''
Created on Jan 27, 2011

@author: oabalbin
'''

import sys
import os
import subprocess
import numpy as np
from optparse import OptionParser

'''
/exds/sw/bioinfo/alignment/samtools/samtools-0.1.10/misc/wgsim -N 1000000 -1 98 -2 98
 -r 0.0010 -R 0.0 ./TP53.fa ./TP53.0.fq ./TP53.1.fq
'''

def simulate_reads(npaired_reads,length, snps_fraction, indels_fraction,
                   error_rate, ref_genome, mate0, mate1, path_to_wgsim):
    '''
    '''
    wgsim_command = os.path.join(path_to_wgsim, "wgsim")
    args=[wgsim_command, '-N', npaired_reads, '-1', length, 
          '-2', length, '-r', snps_fraction, '-R', indels_fraction,
          '-e', error_rate,
          ref_genome, mate0, mate1]
    args=map(str,args)
    subprocess.call(args)    
    pass

def mutate_refGenome(ref_genome, mutation_file=None, outfolder=None):
    '''
    '''
    from Bio import SeqIO
    from Bio.Seq import Seq
    from Bio.SeqRecord import SeqRecord
    from random import sample
    
    # numeber of included mutations
    # new genome
    nmut=10
    nmut_genomes=10
    
    refGenome = open(ref_genome, "rU")
    outfiles_list=[]
    
    if mutation_file is not None:
        mfile = open(mutation_file)
        mutation_list = [line.strip('\n').split('\t') for line in mfile]
                
    for record in SeqIO.parse(refGenome, "fasta") :
        name = record.id.split('_')
        chr, start, end = name[1],int(name[2]), int(name[3])        
        
        mutrecords_file= open(os.path.join(outfolder,
                                     "mutations_record"),'w')
        
        for k in range(nmut_genomes):
            
            # introduce 10 random mutations 
            # from the mutation list. Initially a 15 element list
            random_set=sample(range(len(mutation_list)),nmut)
            #this_mutation_list = [mutation_list[i] for i in random_set]
            this_mutation_list,index=[],[]

            for i in random_set:
                this_mutation_list.append(mutation_list[i])
                index.append(i)
                
            mname=",".join(map(str,index)).replace(',',';')
            
            # Assign original sequence
            mseq =  np.array(record.seq)
            for mut in this_mutation_list:
                
                mchr,mstart,mbase=mut[0],int(mut[1]),mut[2]
                #mname=",".join(mut).replace(',','_')
                if chr == mchr:
                    mpos=mstart-start
                    print mstart, mpos
                    print record.seq[mpos]
                    # Introduce mutation in sequence
                    mseq[mpos]=mbase
                    print mseq[mpos]
            
            mutseq = ",".join(list(mseq)).replace(',','')
            '''
            # to create a list of recors
            mutated_seqs.append(SeqRecord(mutseq, 
                                         id=record.id+mname,
                                         name=mname,
                                         description=mname
                                         ))
            '''
            
            mutrecord = SeqRecord(Seq(mutseq), 
                                         id=record.id+"_"+mname,
                                         name=mname,
                                         description=mname
                                         )
            outfile=os.path.join(outfolder,"TP53_mut%d.fa"%k)
            
            output_handle = open(os.path.join(outfile),'w')
            SeqIO.write(mutrecord, output_handle, "fasta")
            output_handle.close()
            outfiles_list.append(outfile)
            # Eye here
            outline=[str(k),mname]
            mutrecords_file.write(",".join(outline).replace(',','\t')+'\n')
         
            refGenome.close()
    
    return outfiles_list
    

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    
    '''
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="file with run configuration")
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration") 
    '''
    optionparser.add_option("-f", "--mutation_file", dest="mutation_file",
                            help="file with specific mutations")

    optionparser.add_option("-g", "--genome_fasta", dest="genome_fasta",
                            help="file with reference genome")
    optionparser.add_option("-o", "--out_folder", dest="out_folder",
                            help="folder to write the mutated genome")

  
    (options, args) = optionparser.parse_args()
    
    
    # Generate the mutated genomes    
    outfiles_list = mutate_refGenome(options.genome_fasta, options.mutation_file, options.out_folder)
    
    # Simulate reads obtained from your random genome
    # Simulation parameters
    npaired_reads=1e4
    length=98
    snps_fraction=0.0    #Fraction of snps beyond the introduced above
    indels_fraction=0.0  #Fraction of indels
    error_rate=0.0002   #default 0.0
    path_to_wgsim='/exds/sw/bioinfo/alignment/samtools/samtools-0.1.10/misc/'
    
    for i,mutgenome in enumerate(outfiles_list):
        fname = 'mctp_'+mutgenome.split('/')[-1].replace('.fa','')+'_'+str(i)
        folder=mutgenome.split('/')[:-1]
        folder=os.path.join(",".join(folder).replace(',','/'),fname)
        
        os.mkdir(folder)
        mate0=os.path.join(folder,'mate0.fq')
        mate1=os.path.join(folder,'mate1.fq')
        
        simulate_reads(npaired_reads,length, snps_fraction, indels_fraction,
                       error_rate, mutgenome, mate0, mate1, path_to_wgsim)
      



